{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n",
    "\n",
    "Instructions for setting up Colab are as follows:\n",
    "1. Open a new Python 3 notebook.\n",
    "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n",
    "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n",
    "4. Run this cell to set up dependencies.\n",
    "\"\"\"\n",
    "# If you're using Google Colab and not running locally, run this cell.\n",
    "!pip install wget\n",
    "!apt-get install sox\n",
    "!pip install nemo_toolkit[asr]==0.10.0b10\n",
    "!pip install unidecode\n",
    "\n",
    "!mkdir configs\n",
    "!wget -P configs/ https://raw.githubusercontent.com/NVIDIA/NeMo/master/examples/asr/configs/jasper_an4.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction to End-To-End Automatic Speech Recognition\n",
    "\n",
    "This notebook contains a basic tutorial of Automatic Speech Recognition (ASR) concepts, introduced with code snippets using the [NeMo framework](https://github.com/NVIDIA/NeMo).\n",
    "We will first introduce the basics of the main concepts behind speech recognition, then explore concrete examples of what the data looks like and walk through putting together a simple end-to-end ASR pipeline.\n",
    "\n",
    "We assume that you are familiar with general machine learning concepts and can follow Python code, and we'll be using the [AN4 dataset from CMU](http://www.speech.cs.cmu.edu/databases/an4/) (with processing using `sox`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conceptual Overview: What is ASR?\n",
    "\n",
    "ASR, or **Automatic Speech Recognition**, refers to the problem of getting a program to automatically transcribe spoken language (speech-to-text). Our goal is usually to have a model that minimizes the **Word Error Rate (WER)** metric when transcribing speech input. In other words, given some audio file (e.g. a WAV file) containing speech, how do we transform this into the corresponding text with as few errors as possible?\n",
    "\n",
    "Traditional speech recognition takes a generative approach, modeling the full pipeline of how speech sounds are produced: from a **language model** that encapsulates likely orderings of words (e.g. an n-gram model), to a **pronunciation model** for each word in the vocabulary (e.g. a pronunciation table), to an **acoustic model** that translates the pronunciations to audio waveforms (e.g. a Gaussian Mixture Model), and so on.\n",
    "\n",
    "Then, if we receive some spoken input, our goal would be to find the most likely sequence of text that would result in the given audio according to our pipeline of models. Overall, with traditional speech recognition, we try to model `Pr(audio|transcript)*Pr(transcript)`, and take the argmax of this over possible transcripts.\n",
    "\n",
    "Over time, neural nets advanced to the point where each component of the traditional speech recognition model could be replaced by a neural model that had better performance and that had a greater potential for generalization. For example, we could replace an n-gram model with a neural language model, and replace a pronunciation table with a neural pronunciation model, and so on. However, each of these neural models need to be trained individually on different tasks, and errors in any model in the pipeline could throw off the whole prediction.\n",
    "\n",
    "Thus, we can see the appeal of **end-to-end ASR architectures**--discriminative models that simply take an audio input and give a textual output, and in which all components are trained together towards the same goal. A much easier pipeline to handle!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### End-To-End ASR\n",
    "\n",
    "With an end-to-end model, we want to directly learn `Pr(transcript|audio)` in order to predict the transcripts from the original audio. Since we are dealing with sequential information--audio data over time that corresponds to a sequence of letters--RNNs are the obvious choice. But now we have a pressing problem to deal with: since our input sequence (number of audio timesteps) is not the same length as our desired output (transcript length), how do we match each time step from the audio data to the correct output characters?\n",
    "\n",
    "Earlier speech recognition approaches relied on **temporally-aligned data**, in which each segment of time in an audio file was matched up to a corresponding speech sound such as a phoneme or word. However, if we would like to have the flexibility to predict letter-by-letter to prevent OOV (out of vocabulary) issues, then each time step in the data would have to be labeled with the letter sound that the speaker is making at that point in the audio file. With that information, it seems like we should simply be able to try to predict the correct letter for each time step and then collapse the repeated letters (e.g. the prediction output `LLLAAAAPPTOOOPPPP` would become `LAPTOP`). It turns out that this idea has some problems: not only does alignment make the dataset incredibly labor-intensive to label, but also, what do we do with words like \"book\" that contain consecutive repeated letters? Simply squashing repeated letters together would not work in that case!\n",
    "\n",
    "![Alignment example](https://raw.githubusercontent.com/NVIDIA/NeMo/master/examples/asr/notebooks/images/alignment_example.png)\n",
    "\n",
    "Modern end-to-end approaches get around this using methods that don't require manual alignment at all, so that the input-output pairs are really just the raw audio and the transcript--no extra data or labeling required. Let's briefly go over two popular approaches that allow us to do this, Connectionist Temporal Classification (CTC) and sequence-to-sequence models with attention.\n",
    "\n",
    "#### Connectionist Temporal Classification (CTC)\n",
    "\n",
    "In normal speech recognition prediction output, we would expect to have characters such as the letters from A through Z, numbers 0 through 9, spaces (\"\\_\"), and so on. CTC introduces a new intermediate output token called the **blank token** (\"-\") that is useful for getting around the alignment issue.\n",
    "\n",
    "With CTC, we still predict one token per time segment of speech, but we use the blank token to figure out where we can and can't collapse the predictions. The appearance of a blank token helps separate repeating letters that should not be collapsed. For instance, with an audio snippet segmented into `T=11` time steps, we could get predictions that look like `BOO-OOO--KK`, which would then collapse to `\"BO-O-K\"`, and then we would remove the blank tokens to get our final output, `BOOK`.\n",
    "\n",
    "Now, we can predict one output token per time step, then collapse and clean to get sensible output without any fear of ambiguity from repeating letters! A simple way of getting predictions like this would be to apply a bidirectional RNN to the audio input, apply softmax over each time step's output, and then take the token with the highest probability. The method of always taking the best token at each time step is called **greedy decoding, or max decoding**.\n",
    "\n",
    "To calculate our loss for backprop, we would like to know the log probability of the model producing the correct transcript, `log(Pr(transcript|audio))`. We can get the log probability of a single intermediate output sequence (e.g. `BOO-OOO--KK`) by summing over the log probabilities we get from each token's softmax value, but note that the resulting sum is different from the log probability of the transcript itself (`BOOK`). This is because there are multiple possible output sequences of the same length that can be collapsed to get the same transcript (e.g. `BBO--OO-KKK` also results in `BOOK`), and so we need to **marginalize over every valid sequence of length `T` that collapses to the transcript**.\n",
    "\n",
    "Therefore, to get our transcript's log probability given our audio input, we must sum the log probabilities of every sequence of length `T` that collapses to the transcript (e.g. `log(Pr(output: \"BOOK\"|audio)) = log(Pr(BOO-OOO--KK|audio)) + log(Pr(BBO--OO-KKK|audio)) + ...`). In practice, we can use a dynamic programming approach to calculate this, accumulating our log probabilities over different \"paths\" through the softmax outputs at each time step.\n",
    "\n",
    "If you would like a more in-depth explanation of how CTC works, or how we can improve our results by using a modified beam search algorithm, feel free to check out the Further Reading section at the end of this notebook for more resources.\n",
    "\n",
    "#### Sequence-to-Sequence with Attention\n",
    "\n",
    "One problem with CTC is that predictions at different time steps are conditionally independent, which is an issue because the words in a continuous utterance tend to be related to each other in some sensible way. With this conditional independence assumption, we can't learn a language model that can represent such dependencies, though we can add a language model on top of the CTC output to mitigate this to some degree.\n",
    "\n",
    "A popular alternative is to use a sequence-to-sequence model with attention. A typical seq2seq model for ASR consists of some sort of **bidirectional RNN encoder** that consumes the audio sequence timestep-by-timestep, and where the outputs are then passed to an **attention-based decoder**. Each prediction from the decoder is based on attending to some parts of the entire encoded input, as well as the previously outputted tokens.\n",
    "\n",
    "The outputs of the decoder can be anything from word pieces to phonemes to letters, and since predictions are not directly tied to time steps of the input, we can just continue producing tokens one-by-one until an end token is given (or we reach a specified max output length). This way, we do not need to deal with audio alignment, and our predicted transcript is just the sequence of outputs given by our decoder.\n",
    "\n",
    "Now that we have an idea of what some popular end-to-end ASR models look like, let's take a look at the audio data we'll be working with for our example."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Taking a Look at Our Data (AN4)\n",
    "\n",
    "The AN4 dataset, also known as the Alphanumeric dataset, was collected and published by Carnegie Mellon University. It consists of recordings of people spelling out addresses, names, telephone numbers, etc., one letter or number at a time, as well as their corresponding transcripts. We choose to use AN4 for this tutorial because it is relatively small, with 948 training and 130 test utterances, and so it trains quickly.\n",
    "\n",
    "Before we get started, let's download and prepare the dataset. The utterances are available as `.sph` files, so we will need to convert them to `.wav` for later processing. Please make sure you have [Sox](http://sox.sourceforge.net/) installed for this step (see the \"Downloads\" section of the main page)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This is where the an4/ directory will be placed.\n",
    "# Change this if you don't want the data to be extracted in the current directory.\n",
    "data_dir = '.'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import os\n",
    "import subprocess\n",
    "import tarfile\n",
    "import wget\n",
    "\n",
    "# Download the dataset. This will take a few moments...\n",
    "print(\"******\")\n",
    "if not os.path.exists(data_dir + '/an4_sphere.tar.gz'):\n",
    "    an4_url = 'http://www.speech.cs.cmu.edu/databases/an4/an4_sphere.tar.gz'\n",
    "    an4_path = wget.download(an4_url, data_dir)\n",
    "    print(f\"Dataset downloaded at: {an4_path}\")\n",
    "else:\n",
    "    print(\"Tarfile already exists.\")\n",
    "    an4_path = data_dir + '/an4_sphere.tar.gz'\n",
    "\n",
    "# Untar and convert .sph to .wav (using sox)\n",
    "tar = tarfile.open(an4_path)\n",
    "tar.extractall(path=data_dir)\n",
    "\n",
    "print(\"Converting .sph to .wav...\")\n",
    "sph_list = glob.glob(data_dir + '/an4/**/*.sph', recursive=True)\n",
    "for sph_path in sph_list:\n",
    "    wav_path = sph_path[:-4] + '.wav'\n",
    "    cmd = [\"sox\", sph_path, wav_path]\n",
    "    subprocess.run(cmd)\n",
    "print(\"Finished conversion.\\n******\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You should now have a folder called `an4` that contains `etc/an4_train.transcription`, `etc/an4_test.transcription`, audio files in `wav/an4_clstk` and `wav/an4test_clstk`, along with some other files we will not need.\n",
    "\n",
    "Now we can load and take a look at the data. As an example, file `cen2-mgah-b.wav` is a 2.6 second-long audio recording of a man saying the letters \"G L E N N\" one-by-one (feel free to check this out by listening to `./an4/wav/an4_clstk/mgah/cen2-mgah-b.wav`). In an ASR task, the WAV file would be our input, and \"G L E N N\" would be our desired output.\n",
    "\n",
    "Let's plot the waveform, which is simply a line plot of the sequence of values that we read from the file. This is a format of viewing audio that you are likely to be familiar with seeing in many audio editors and visualizers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import librosa\n",
    "import librosa.display\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "# Plot our example audio file's waveform\n",
    "example_file = data_dir + '/an4/wav/an4_clstk/mgah/cen2-mgah-b.wav'\n",
    "audio, sample_rate = librosa.load(example_file)\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (15,7)\n",
    "plt.title('Waveform of Audio Example')\n",
    "plt.ylabel('Amplitude')\n",
    "\n",
    "_ = librosa.display.waveplot(audio)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see the activity in the waveform that corresponds to each letter in the audio, as our speaker here enunciates quite clearly!\n",
    "You can kind of tell that each spoken letter has a different \"shape,\" and it's interesting to note that last two blobs look relatively similar, which is expected because they are both the letter \"N.\"\n",
    "\n",
    "### Spectrograms and Mel Spectrograms\n",
    "\n",
    "However, since audio information is more useful in the context of frequencies of sound over time, we can get a better representation than this raw sequence of 57,330 values.\n",
    "We can apply a [Fourier Transform](https://en.wikipedia.org/wiki/Fourier_transform) on our audio signal to get something more useful: a **spectrogram**, which is a representation of the energy levels (i.e. amplitude, or \"loudness\") of each frequency (i.e. pitch) of the signal over the duration of the file.\n",
    "A spectrogram (which can be viewed as a heat map) is a good way of seeing how the *strengths of various frequencies in the audio vary over time*, and is obtained by breaking up the signal into smaller, usually overlapping chunks and performing a Short-Time Fourier Transform (STFT) on each.\n",
    "\n",
    "Let's examine what the spectrogram of our sample looks like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get spectrogram using Librosa's Short-Time Fourier Transform (stft)\n",
    "spec = np.abs(librosa.stft(audio))\n",
    "spec_db = librosa.amplitude_to_db(spec, ref=np.max)  # Decibels\n",
    "\n",
    "# Use log scale to view frequencies\n",
    "librosa.display.specshow(spec_db, y_axis='log', x_axis='time')\n",
    "plt.colorbar()\n",
    "plt.title('Audio Spectrogram');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, we are able to see each letter being pronounced, and that the last two blobs that correspond to the \"N\"s are pretty similar-looking. But how do we interpret these shapes and colors? Just as in the waveform plot before, we see time passing on the x-axis (all 2.6s of audio). But now, the y-axis represents different frequencies (on a log scale), and *the color on the plot shows the strength of a frequency at a particular point in time*.\n",
    "\n",
    "We're still not done yet, as we can make one more potentially useful tweak: using the **Mel Spectrogram** instead of the normal spectrogram. This is simply a change in the frequency scale that we use from linear (or logarithmic) to the mel scale, which is \"a perceptual scale of pitches judged by listeners to be equal in distance from one another\" (from [Wikipedia](https://en.wikipedia.org/wiki/Mel_scale)).\n",
    "\n",
    "In other words, it's a transformation of the frequencies to be more aligned to what humans perceive; a change of +1000Hz from 2000Hz->3000Hz sounds like a larger difference to us than 9000Hz->10000Hz does, so the mel scale normalizes this such that equal distances sound like equal differences to the human ear. Intuitively, we use the mel spectrogram because in this case we are processing and transcribing human speech, such that transforming the scale to better match what we hear is a useful procedure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the mel spectrogram of our sample\n",
    "mel_spec = librosa.feature.melspectrogram(audio, sr=sample_rate)\n",
    "mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)\n",
    "\n",
    "librosa.display.specshow(\n",
    "    mel_spec_db, x_axis='time', y_axis='mel')\n",
    "plt.colorbar()\n",
    "plt.title('Mel Spectrogram');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building a Simple ASR Pipeline in NeMo\n",
    "\n",
    "Now that we have an idea of what the audio data looks like, we can start building our end-to-end ASR pipeline!\n",
    "\n",
    "We'll be using the **Neural Modules (NeMo) toolkit** for this part, so if you haven't already, you should download and install NeMo and its dependencies. To do so, just follow the directions on the [GitHub page](https://github.com/NVIDIA/NeMo), or in the [documentation](https://nvidia.github.io/NeMo/).\n",
    "\n",
    "NeMo lets us easily hook together the components (modules) of our model, such as the data layer, intermediate layers, and various losses, without worrying too much about implementation details of individual parts or connections between modules. If you're curious, you can read more about NeMo and how it works in the documentation pages linked above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# NeMo's \"core\" package\n",
    "import nemo\n",
    "# NeMo's ASR collection\n",
    "import nemo.collections.asr as nemo_asr"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Creating Data Manifests\n",
    "\n",
    "The first thing we need to do now is to create manifests for our training and evaluation data, which will contain the metadata of our audio files. NeMo data layers take in a standardized manifest format where each line corresponds to one sample of audio, such that the number of lines in a manifest is equal to the number of samples that are represented by that manifest. A line must contain the path to an audio file, the corresponding transcript (or path to a transcript file), and the duration of the audio sample.\n",
    "\n",
    "Here's an example of what one line in a NeMo-compatible manifest might look like:\n",
    "```\n",
    "{\"audio_filepath\": \"path/to/audio.wav\", \"duration\": 3.45, \"text\": \"this is a nemo tutorial\"}\n",
    "```\n",
    "\n",
    "We can build our training and evaluation manifests using `an4/etc/an4_train.transcription` and `an4/etc/an4_test.transcription`, which have lines containing transcripts and their corresponding audio file IDs:\n",
    "```\n",
    "...\n",
    "<s> P I T T S B U R G H </s> (cen5-fash-b)\n",
    "<s> TWO SIX EIGHT FOUR FOUR ONE EIGHT </s> (cen7-fash-b)\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Building Manifest Files --- #\n",
    "import json\n",
    "\n",
    "# Function to build a manifest\n",
    "def build_manifest(transcripts_path, manifest_path, wav_path):\n",
    "    with open(transcripts_path, 'r') as fin:\n",
    "        with open(manifest_path, 'w') as fout:\n",
    "            for line in fin:\n",
    "                # Lines look like this:\n",
    "                # <s> transcript </s> (fileID)\n",
    "                transcript = line[: line.find('(')-1].lower()\n",
    "                transcript = transcript.replace('<s>', '').replace('</s>', '')\n",
    "                transcript = transcript.strip()\n",
    "\n",
    "                file_id = line[line.find('(')+1 : -2]  # e.g. \"cen4-fash-b\"\n",
    "                audio_path = os.path.join(\n",
    "                    data_dir, wav_path,\n",
    "                    file_id[file_id.find('-')+1 : file_id.rfind('-')],\n",
    "                    file_id + '.wav')\n",
    "\n",
    "                duration = librosa.core.get_duration(filename=audio_path)\n",
    "\n",
    "                # Write the metadata to the manifest\n",
    "                metadata = {\n",
    "                    \"audio_filepath\": audio_path,\n",
    "                    \"duration\": duration,\n",
    "                    \"text\": transcript\n",
    "                }\n",
    "                json.dump(metadata, fout)\n",
    "                fout.write('\\n')\n",
    "                \n",
    "# Building Manifests\n",
    "print(\"******\")\n",
    "train_transcripts = data_dir + '/an4/etc/an4_train.transcription'\n",
    "train_manifest = data_dir + '/an4/train_manifest.json'\n",
    "build_manifest(train_transcripts, train_manifest, 'an4/wav/an4_clstk')\n",
    "print(\"Training manifest created.\")\n",
    "\n",
    "test_transcripts = data_dir + '/an4/etc/an4_test.transcription'\n",
    "test_manifest = data_dir + '/an4/test_manifest.json'\n",
    "build_manifest(test_transcripts, test_manifest, 'an4/wav/an4test_clstk')\n",
    "print(\"Test manifest created.\")\n",
    "print(\"******\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building Training and Evaluation DAGs\n",
    "\n",
    "Let's take a look at the model that we will be building, and how we specify its parameters.\n",
    "\n",
    "#### The Jasper Model\n",
    "\n",
    "We will be putting together a small [Jasper (Just Another SPeech Recognizer) model](https://arxiv.org/abs/1904.03288).\n",
    "In brief, Jasper architectures consist of a repeated block structure that utilizes 1D convolutions.\n",
    "In a Jasper_KxR model, `R` sub-blocks (consisting of a 1D convolution, batch norm, ReLU, and dropout) are grouped into a single block, which is then repeated `K` times.\n",
    "We also have a one extra block at the beginning and a few more at the end that are invariant of `K` and `R`, and we use CTC loss.\n",
    "\n",
    "A Jasper model looks like roughly this:\n",
    "\n",
    "![Jasper with CTC](https://raw.githubusercontent.com/NVIDIA/NeMo/master/docs/sources/source/asr/jasper_vertical.png)\n",
    "\n",
    "#### Specifying Our Model with a YAML Config File\n",
    "\n",
    "For this tutorial, we'll build a *Jasper_4x1 model*, with `K=4` blocks of single (`R=1`) sub-blocks and a *greedy CTC decoder*, using the configuration found in `./configs/jasper_an4.yaml`.\n",
    "\n",
    "If we open up this config file, we find that there is an entry labeled `JasperEncoder`, with a field called `jasper` that contains a list with multiple entries. Each of the members in this list specifies one block in our model, and looks something like this:\n",
    "```\n",
    "- filters: 256\n",
    "  repeat: 1\n",
    "  kernel: [11]\n",
    "  stride: [2]\n",
    "  dilation: [1]\n",
    "  dropout: 0.2\n",
    "  residual: false\n",
    "```\n",
    "The first member of the list corresponds to the first block in the Jasper architecture diagram, which appears regardless of `K` and `R`.\n",
    "Next, we have four entries that correspond to the `K=4` blocks, and each has `repeat: 1` since we are using `R=1`.\n",
    "These are followed by two more entries for the blocks that appear at the end of our Jasper model before the CTC loss.\n",
    "\n",
    "There are also some entries at the top of the file that specify that we should be shuffling our training data but not our evaluation data (see `AudioToTextDataLayer_train` and `AudioToTextDataLayer_eval`), and some specifications for preprocessing and converting the audio data (in `AudioToMelSpectrogramPreprocessor`).\n",
    "\n",
    "Using a YAML config such as this is helpful for getting a quick and human-readable overview of what your architecture looks like, and allows you to swap out model and run configurations easily without needing to change your code.\n",
    "\n",
    "#### Building Training and Evaluation DAGs with NeMo\n",
    "\n",
    "Building a model using NeMo consists of (1) instantiating the neural modules we need and (2) specifying the DAG by linking them together."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In NeMo, **the training and inference pipelines are managed by a `NeuralModuleFactory`**, which takes care of checkpointing, callbacks, and logs, along with other details in training and inference. We set its `log_dir` argument to specify where our model logs and outputs will be written, and can set other training and inference settings in its constructor. For instance, if we were **resuming training from a checkpoint**, we would set the argument `checkpoint_dir=<path_to_checkpoint>`.\n",
    "\n",
    "Along with logs in NeMo, you can optionally view the tensorboard logs with the `create_tb_writer=True` argument to the `NeuralModuleFactory`. By default all the tensorboard log files will be stored in `{log_dir}/tensorboard`, but you can change this with the `tensorboard_dir` argument. One can load tensorboard logs through tensorboard by running `tensorboard --logdir=<path_to_tensorboard dir>` in the terminal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create our NeuralModuleFactory, which will oversee the neural modules.\n",
    "neural_factory = nemo.core.NeuralModuleFactory(\n",
    "    log_dir=data_dir+'/an4_tutorial/',\n",
    "    create_tb_writer=True)\n",
    "\n",
    "logger = nemo.logging"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that if you would like your pipeline to support running on CPU, you should\n",
    "instead add these imports and then add an additional 'placement' argument:\n",
    "\n",
    "```python\n",
    "import torch\n",
    "from nemo.core import DeviceType\n",
    "\n",
    "neural_factory = nemo.core.NeuralModuleFactory(\n",
    "    log_dir=data_dir+'/an4_tutorial/',\n",
    "    placement=(DeviceType.GPU if torch.cuda.is_available() else DeviceType.CPU)\n",
    "    )\n",
    "```\n",
    "\n",
    "Now that we have our neural module factory, we can **specify our neural modules and instantiate them**. Here, we load the parameters for each module from the configuration file using `import_from_config`. For the module parameters that we can't read directly from the config file or want to overwrite, we can use the `overwrite_params` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Config Information ---#\n",
    "from ruamel.yaml import YAML\n",
    "\n",
    "config_path = '../configs/jasper_an4.yaml'\n",
    "\n",
    "yaml = YAML(typ='safe')\n",
    "with open(config_path) as f:\n",
    "    params = yaml.load(f)\n",
    "labels = params['labels'] # Vocab\n",
    "\n",
    "# --- Instantiate Neural Modules --- #\n",
    "\n",
    "# Create training and test data layers (which load data) and data preprocessor\n",
    "data_layer_train = nemo_asr.AudioToTextDataLayer.import_from_config(\n",
    "    config_path,\n",
    "    \"AudioToTextDataLayer_train\",\n",
    "    overwrite_params={\"manifest_filepath\": train_manifest}\n",
    ") # Training datalayer\n",
    "\n",
    "data_layer_test = nemo_asr.AudioToTextDataLayer.import_from_config(\n",
    "    config_path,\n",
    "    \"AudioToTextDataLayer_eval\",\n",
    "    overwrite_params={\"manifest_filepath\": test_manifest}\n",
    ") # Eval datalayer\n",
    "\n",
    "data_preprocessor = nemo_asr.AudioToMelSpectrogramPreprocessor.import_from_config(\n",
    "    config_path, \"AudioToMelSpectrogramPreprocessor\"\n",
    ")\n",
    "\n",
    "# Create the Jasper_4x1 encoder as specified, and a CTC decoder\n",
    "encoder = nemo_asr.JasperEncoder.import_from_config(\n",
    "    config_path, \"JasperEncoder\"\n",
    ")\n",
    "\n",
    "decoder = nemo_asr.JasperDecoderForCTC.import_from_config(\n",
    "    config_path, \"JasperDecoderForCTC\",\n",
    "    overwrite_params={\"num_classes\": len(labels)}\n",
    ")\n",
    "\n",
    "ctc_loss = nemo_asr.CTCLossNM(num_classes=len(labels))\n",
    "greedy_decoder = nemo_asr.GreedyCTCDecoder()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next step is to assemble our training DAG by specifying the inputs to each neural module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Assemble Training DAG --- #\n",
    "audio_signal, audio_signal_len, transcript, transcript_len = data_layer_train()\n",
    "\n",
    "processed_signal, processed_signal_len = data_preprocessor(\n",
    "    input_signal=audio_signal,\n",
    "    length=audio_signal_len)\n",
    "\n",
    "encoded, encoded_len = encoder(\n",
    "    audio_signal=processed_signal,\n",
    "    length=processed_signal_len)\n",
    "\n",
    "log_probs = decoder(encoder_output=encoded)\n",
    "preds = greedy_decoder(log_probs=log_probs)  # Training predictions\n",
    "loss = ctc_loss(\n",
    "    log_probs=log_probs,\n",
    "    targets=transcript,\n",
    "    input_length=encoded_len,\n",
    "    target_length=transcript_len)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We would like to be able to evaluate our model on the test set, as well, so let's set up the evaluation DAG.\n",
    "\n",
    "Our **evaluation DAG will reuse most of the parts of the training DAG with the exception of the data layer**, since we are loading the evaluation data from a different file but evaluating on the same model. Note that if we were using data augmentation in training, we would also leave that out in the evaluation DAG."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Assemble Validation DAG --- #\n",
    "(audio_signal_test, audio_len_test,\n",
    " transcript_test, transcript_len_test) = data_layer_test()\n",
    "\n",
    "processed_signal_test, processed_len_test = data_preprocessor(\n",
    "    input_signal=audio_signal_test,\n",
    "    length=audio_len_test)\n",
    "\n",
    "encoded_test, encoded_len_test = encoder(\n",
    "    audio_signal=processed_signal_test,\n",
    "    length=processed_len_test)\n",
    "\n",
    "log_probs_test = decoder(encoder_output=encoded_test)\n",
    "preds_test = greedy_decoder(log_probs=log_probs_test)  # Test predictions\n",
    "loss_test = ctc_loss(\n",
    "    log_probs=log_probs_test,\n",
    "    targets=transcript_test,\n",
    "    input_length=encoded_len_test,\n",
    "    target_length=transcript_len_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running the Model\n",
    "\n",
    "We would like to be able to monitor our model while it's training, so we use **callbacks**. In general, *callbacks are functions that are called at specific intervals over the course of training or inference*, such as at the start or end of every *n* iterations, epochs, etc. The callbacks we'll be using for this are the `SimpleLossLoggerCallback`, which reports the training loss (or another metric of your choosing, such as WER for ASR tasks), and the `EvaluatorCallback`, which regularly evaluates the model on the test set. Both of these callbacks require you to pass in the tensors to be evaluated--these would be the final outputs of the training and eval DAGs above.\n",
    "\n",
    "Another useful callback is the `CheckpointCallback`, for saving checkpoints at set intervals. We create one here just to demonstrate how it works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Create Callbacks --- #\n",
    "\n",
    "# We use these imports to pass to callbacks more complex functions to perform.\n",
    "from nemo.collections.asr.helpers import monitor_asr_train_progress, \\\n",
    "    process_evaluation_batch, process_evaluation_epoch\n",
    "from functools import partial\n",
    "\n",
    "train_callback = nemo.core.SimpleLossLoggerCallback(\n",
    "    # Notice that we pass in loss, predictions, and the transcript info.\n",
    "    # Of course we would like to see our training loss, but we need the\n",
    "    # other arguments to calculate the WER.\n",
    "    tensors=[loss, preds, transcript, transcript_len],\n",
    "    # The print_func defines what gets printed.\n",
    "    print_func=partial(\n",
    "        monitor_asr_train_progress,\n",
    "        labels=labels),\n",
    "    tb_writer=neural_factory.tb_writer\n",
    "    )\n",
    "\n",
    "# We can create as many evaluation DAGs and callbacks as we want,\n",
    "# which is useful in the case of having more than one evaluation dataset.\n",
    "# In this case, we only have one.\n",
    "eval_callback = nemo.core.EvaluatorCallback(\n",
    "    eval_tensors=[loss_test, preds_test, transcript_test, transcript_len_test],\n",
    "    user_iter_callback=partial(\n",
    "        process_evaluation_batch, labels=labels),\n",
    "    user_epochs_done_callback=process_evaluation_epoch,\n",
    "    eval_step=500,  # How often we evaluate the model on the test set\n",
    "    tb_writer=neural_factory.tb_writer\n",
    "    )\n",
    "\n",
    "checkpoint_saver_callback = nemo.core.CheckpointCallback(\n",
    "    folder=data_dir+'/an4_checkpoints',\n",
    "    step_freq=1000  # How often checkpoints are saved\n",
    "    )\n",
    "\n",
    "if not os.path.exists(data_dir+'/an4_checkpoints'):\n",
    "    os.makedirs(data_dir+'/an4_checkpoints')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our model and callbacks set up, **how do we run it**?\n",
    "\n",
    "Once we create our neural factory and the callbacks for the information that we want to see, we can **start training** by simply calling the train function on the tensors we want to optimize and our callbacks!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# --- Start Training! --- #\n",
    "neural_factory.train(\n",
    "    tensors_to_optimize=[loss],\n",
    "    callbacks=[train_callback, eval_callback, checkpoint_saver_callback],\n",
    "    optimizer='novograd',\n",
    "    optimization_params={\n",
    "        \"num_epochs\": 100, \"lr\": 0.01, \"weight_decay\": 1e-4\n",
    "    })\n",
    "\n",
    "# Training for 100 epochs will take a little while, depending on your machine.\n",
    "# It should take about 20 minutes on Google Colab.\n",
    "# At the end of 100 epochs, your evaluation WER should be around 20-25%."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There we go! We've put together a full training pipeline for the model and trained it for 100 epochs.\n",
    "\n",
    "### Inference\n",
    "\n",
    "What if we have a trained model that we **just want to run inference** on?\n",
    "\n",
    "In that case, we just need to instantiate and link up the modules needed for the evaluation DAG (same procedure as before), and then run `infer` to get the results. Let's see what performing inference with our last checkpoint would look like."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# --- Inference Only --- #\n",
    "\n",
    "# We've already built the inference DAG above, so all we need is to call infer().\n",
    "evaluated_tensors = neural_factory.infer(\n",
    "    # These are the tensors we want to get from the model.\n",
    "    tensors=[loss_test, preds_test, transcript_test, transcript_len_test],\n",
    "    # checkpoint_dir specifies where the model params are loaded from.\n",
    "    checkpoint_dir=(data_dir+'/an4_checkpoints')\n",
    "    )\n",
    "\n",
    "# Process the results to get WER\n",
    "from nemo.collections.asr.helpers import word_error_rate, \\\n",
    "    post_process_predictions, post_process_transcripts\n",
    "\n",
    "greedy_hypotheses = post_process_predictions(\n",
    "    evaluated_tensors[1], labels)\n",
    "\n",
    "references = post_process_transcripts(\n",
    "    evaluated_tensors[2], evaluated_tensors[3], labels)\n",
    "\n",
    "wer = word_error_rate(hypotheses=greedy_hypotheses, references=references)\n",
    "print(\"*** Greedy WER: {:.2f} ***\".format(wer * 100))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And that's it!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Improvements\n",
    "\n",
    "In this section, we'll briefly cover a few possibilities for improving an ASR model.\n",
    "\n",
    "### Data Augmentation\n",
    "\n",
    "There exist several ASR data augmentation methods that can increase the size of our training set.\n",
    "\n",
    "For example, we can perform augmentation on the spectrograms by zeroing out specific frequency segments (\"frequency masking\") or time segments (\"time masking\") as described by [SpecAugment](https://arxiv.org/abs/1904.08779), or zero out rectangles on the spectrogram as in [Cutout](https://arxiv.org/pdf/1708.04552.pdf). In NeMo, we can do all three of these by simply adding in a `SpectrogramAugmentation` neural module. (As of now, it does not perform the time warping from the SpecAugment paper.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a SpectrogramAugmentation module\n",
    "spectrogram_aug = nemo_asr.SpectrogramAugmentation(\n",
    "    rect_masks=5, rect_time=120, rect_freq=50)\n",
    "\n",
    "# Rearrange training DAG to use augmentation.\n",
    "# The following code is mostly copy/pasted from the \"Assemble Training DAG\"\n",
    "# section, with only one line added!\n",
    "audio_signal, audio_signal_len, transcript, transcript_len = data_layer_train()\n",
    "\n",
    "processed_signal, processed_signal_len = data_preprocessor(\n",
    "    input_signal=audio_signal,\n",
    "    length=audio_signal_len)\n",
    "\n",
    "############## This is the only part that's changed! ##############\n",
    "processed_signal_aug = spectrogram_aug(input_spec=processed_signal)\n",
    "\n",
    "encoded, encoded_len = encoder(\n",
    "    audio_signal=processed_signal_aug,  # Change this argument too\n",
    "    length=processed_signal_len)\n",
    "###################################################################\n",
    "\n",
    "log_probs = decoder(encoder_output=encoded)\n",
    "preds = greedy_decoder(log_probs=log_probs)  # Training predictions\n",
    "loss = ctc_loss(\n",
    "    log_probs=log_probs,\n",
    "    targets=transcript,\n",
    "    input_length=encoded_len,\n",
    "    target_length=transcript_len)\n",
    "\n",
    "# And then you can train as usual.\n",
    "# If you want to try it out in this notebook,\n",
    "# be sure to run neural_factory.reset_trainer() right before training again."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Another popular method of ASR data augmentation is speed perturbation, where the audio is sped up or slowed down slightly (e.g. 10% faster or slower), but that's currently not available as an option in NeMo yet.\n",
    "\n",
    "### Using a Language Model\n",
    "\n",
    "Though a language model may not be especially suited to a task like AN4 where we have a bunch of letters being read in sequence, adding a language model for inference can the improve WER in most other ASR tasks, where the speech more closely matches normal patterns. We can use the probability distribution that a language model gives us to better match our predictions to sequences of words we would be more likely to see in the real world, such as correcting \"keyboard and house\" to \"keyboard and mouse.\"\n",
    "\n",
    "If you have a language model that you'd like to use with a NeMo model, you can add a `BeamSearchDecoderWithLM` module to your DAG to get beam search predictions that use your language model file.\n",
    "\n",
    "```python\n",
    "### Instantiating the module ###\n",
    "beam_search_lm = nemo_asr.BeamSearchDecoderWithLM(\n",
    "    vocab=labels, beam_width=16,\n",
    "    alpha=2, beta=1.5,\n",
    "    lm_path=\"<path/to/language/model>\"\n",
    ")\n",
    "...\n",
    "\n",
    "### In the DAG ###\n",
    "beam_predictions = beam_search_lm(log_probs=log_probs, lob_probs_len=encoded_len)\n",
    "...\n",
    "\n",
    "### During inference ###\n",
    "evaluated_tensors = neural_factory.infer(\n",
    "    tensors=[loss_test, preds_test, transcript_test, transcript_len_test, beam_predictions],\n",
    "    checkpoint_dir=(data_dir+'/an4_checkpoints')\n",
    ")\n",
    "\n",
    "# Get the top beam search hypothesis for each sample\n",
    "beam_hypotheses = [] \n",
    "for mini_batch in evaluated_tensors[-1]:\n",
    "    for sample_hypotheses in mini_batch:\n",
    "        # sample_hypotheses is a set of (probability, prediction) pairs\n",
    "        beam_hypotheses.append(sample_hypotheses[0][1]) # Take top prediction\n",
    "\n",
    "references = post_process_transcripts(evaluated_tensors[1], labels=labels, \n",
    "                                      transcript_len_list=evaluated_tensors[2])\n",
    "wer = word_error_rate(hypotheses=beam_hypotheses, references=references)\n",
    "\n",
    "print(\"BEAM WER {:.2f}\".format(wer*100))\n",
    "```\n",
    "\n",
    "### Fast Training\n",
    "\n",
    "Last but not least, we could simply speed up training our model! If you have the resources, you can speed up training by splitting the workload across multiple GPUs. Otherwise (or in addition), there's always mixed precision training, which allows you to increase your batch size.\n",
    "\n",
    "You can read more about both mixed precision and multi-GPU training using NeMo on [this page of the documentation](https://nvidia.github.io/NeMo/training.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Further reading/watching:\n",
    "\n",
    "That's all for now! If you'd like to learn more about the topics covered in this tutorial, here are some resources that may interest you:\n",
    "- [Stanford Lecture on ASR](https://www.youtube.com/watch?v=3MjIkWxXigM)\n",
    "- [\"An Intuitive Explanation of Connectionist Temporal Classification\"](https://towardsdatascience.com/intuitively-understanding-connectionist-temporal-classification-3797e43a86c)\n",
    "- [Explanation of CTC with Prefix Beam Search](https://medium.com/corti-ai/ctc-networks-and-language-models-prefix-beam-search-explained-c11d1ee23306)\n",
    "- [Listen Attend and Spell Paper (seq2seq ASR model)](https://arxiv.org/abs/1508.01211)\n",
    "- [Explanation of the mel spectrogram in more depth](https://towardsdatascience.com/getting-to-know-the-mel-spectrogram-31bca3e2d9d0)\n",
    "- [Jasper Paper](https://arxiv.org/abs/1904.03288)\n",
    "- [SpecAugment Paper](https://arxiv.org/abs/1904.08779)\n",
    "- [Explanation and visualization of SpecAugment](https://towardsdatascience.com/state-of-the-art-audio-data-augmentation-with-google-brains-specaugment-and-pytorch-d3d1a3ce291e)\n",
    "- [Cutout Paper](https://arxiv.org/pdf/1708.04552.pdf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
